#ifndef __ROLE_HPP__
#define __ROLE_HPP__

#include "utils.hpp"
#include "malloc.h"
#include <sched.h>
#include <sys/time.h>
typedef long long ll;
#define KB (1<<10)
#define MB (KB<<10)
#define GB (MB<<10)
#define MR_PERMISSION                                                  \
  (IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ | \
   IBV_ACCESS_REMOTE_ATOMIC)

#define BUF_FACT 128LL
#define POLL_NUM 16
#define SEND_SIG_INTERVAL 100

volatile bool poll_flag; //记录是否完成poll


struct ibv_mr *memory_alloc(ib_inf *inf, char *name, ll siz){
  assert(sizeof(void*) == sizeof(ll));
  void *mem = memalign(4*KB, siz);
  memset(mem, 0, siz);
  assert(mem);
  struct ibv_mr *mem_mr = ibv_reg_mr(inf->pd, mem, siz, MR_PERMISSION);
  assert(mem_mr);
  ib_mr_attr attr;
  attr.addr = (uint64_t)mem_mr->addr;
  attr.rkey = mem_mr->rkey;
  mt_set(inf->mi, name, (char *)&attr, sizeof(attr));
  return mem_mr;
}
int post_recv(uint32_t size, uint64_t id, struct ibv_qp *qp, uint64_t loc_buf,
              uint32_t lkey){
  struct ibv_sge li = {loc_buf, size, lkey};
  struct ibv_recv_wr *bad_wr, recv_wr;
  memset(&recv_wr, 0, sizeof(recv_wr));
  recv_wr.wr_id = id;
  recv_wr.sg_list = &li;
  recv_wr.num_sge = 1;
  return ibv_post_recv(qp, &recv_wr, &bad_wr);
}

void run_server(struct input_params *conf){
  int res;
  conf->role_int = ROLE_SERVER;
  ib_inf *node_inf = ib_complete_setup(conf);
  int *test_mode, *blk_size;
  while(mt_get(node_inf->mi, "test_mode", (char **)&test_mode, NULL))usleep(1000);
  while(mt_get(node_inf->mi, "blk_size", (char **)&blk_size, NULL))usleep(1000);

  cout<<"test_mode is "<<*test_mode<<endl;
  cout<<"block size is "<<*blk_size<<endl;
  char name[STR_LEN];
  ibv_mr **loc_mr = (ibv_mr **)malloc(conf->num_qps * sizeof(ibv_mr *));
  assert(loc_mr);
  for(int i = 0; i < conf->num_qps; i++){
    sprintf(name, "mem-%d", i);
    loc_mr[i] = memory_alloc(node_inf, name, *blk_size * BUF_FACT);
  }

  if(*test_mode == TEST_WRITE || *test_mode == TEST_READ){
    mt_set(node_inf->mi, "ready", "ready", strlen("ready"));
    //for(int i= 0;i<conf->num_qps;i++)cout<<string((char*)loc_mr[i]->addr)<<endl;
    //do nothing
    for(int i = 0; i < conf->num_qps; i++)
      UNUSED(loc_mr[i]);
  }else if(*test_mode == TEST_SEND){
    queue<uint64_t> *buf_que = new queue<uint64_t> [conf->num_qps];
    uint64_t real_buf;
    for(int i = 0; i < conf->num_qps; i++){
      for(int j = 0; j < BUF_FACT; j++){
        real_buf = (uint64_t)loc_mr[i]->addr + j * conf->blk_size;
        res = post_recv(conf->blk_size, i, node_inf->local_qp[i], 
                  real_buf, loc_mr[i]->lkey);
        assert(res == 0);
        buf_que[i].push(real_buf);
      }
    }
    mt_set(node_inf->mi, "ready", "ready", strlen("ready"));
    ll idx = 0, tot_cnt = conf->iter_num * conf->num_qps, poll_cnt;
    ibv_wc *wc = (ibv_wc *)malloc(POLL_NUM * sizeof(ibv_wc));
    while(idx < tot_cnt){
      poll_cnt = ibv_poll_cq(node_inf->recv_cq, POLL_NUM, wc);
      for(int i = 0; i < poll_cnt; i++){
        if(wc[i].status){cout<<"poll wrong"<<endl;assert(0);}
        uint64_t wr_id = wc[i].wr_id, tb = buf_que[wr_id].front();
        //buf_que[wr_id].push(tb);
        //buf_que[wc[i].wr_id].pop();
        res = post_recv(conf->blk_size, wr_id, node_inf->local_qp[wr_id],
                  tb, loc_mr[wr_id]->lkey);
        assert(res == 0);
      }
      idx += poll_cnt;
    }
  }
  while(mt_get(node_inf->mi, "done", NULL, NULL))usleep(1000);
}

void *poll_routine(void *parm){
  int ret;
  cpu_set_t cpu_mask;
  CPU_ZERO(&cpu_mask);
  CPU_SET(11, &cpu_mask);
  ret = pthread_setaffinity_np(pthread_self() ,sizeof(cpu_set_t), &cpu_mask);
  assert(ret == 0);

  ib_inf *inf = (ib_inf *)parm;
  struct ibv_wc *wc = (struct ibv_wc *)malloc(POLL_NUM * sizeof(struct ibv_wc));
  assert(wc);
  struct ibv_cq *cq = inf->send_cq;
  ll i = 1, tot_rep = inf->inp->iter_num * inf->inp->num_qps;
  int n, j;
  while(i <= tot_rep){
    n = ibv_poll_cq(cq, POLL_NUM, wc);
    for(j = 0; j < n; j++){
      if(wc[j].status){
        printf("poll wrong status %d  id %d n %d\n",wc[j].status, (int)wc[j].wr_id, n);
        break;
      }
    }
    i += n * SEND_SIG_INTERVAL;
  }
  poll_flag = true;
  cout<<"poll thread finish"<<endl;
  return NULL;
}
int post_write(uint32_t size, uint64_t id, struct ibv_qp *qp, uint64_t loc_buf,
              uint32_t lkey, uint64_t rmt_buf, uint32_t rkey, bool sig){
  struct ibv_sge li = {loc_buf, size, lkey};
  struct ibv_send_wr *bad_wr, send_wr;
  memset(&send_wr, 0, sizeof(send_wr));
  send_wr.wr_id = id;
  send_wr.sg_list = &li;
  send_wr.num_sge = 1;
  send_wr.opcode = IBV_WR_RDMA_WRITE;
  if(sig)send_wr.send_flags = IBV_SEND_SIGNALED;
  send_wr.wr.rdma.remote_addr = rmt_buf;
  send_wr.wr.rdma.rkey = rkey;
  return ibv_post_send(qp, &send_wr, &bad_wr);
}
int post_read(uint32_t size, uint64_t id, struct ibv_qp *qp, uint64_t loc_buf,
              uint32_t lkey, uint64_t rmt_buf, uint32_t rkey, bool flg){
  struct ibv_sge li = {loc_buf, size, lkey};
  struct ibv_send_wr *bad_wr, send_wr;
  memset(&send_wr, 0, sizeof(send_wr));
  send_wr.wr_id = id;
  send_wr.sg_list = &li;
  send_wr.num_sge = 1;
  send_wr.opcode = IBV_WR_RDMA_READ;
  if(flg)send_wr.send_flags = IBV_SEND_SIGNALED;
  send_wr.wr.rdma.remote_addr = rmt_buf;
  send_wr.wr.rdma.rkey = rkey;
  return ibv_post_send(qp, &send_wr, &bad_wr);
}
int post_send(uint32_t size, uint64_t id, struct ibv_qp *qp, uint64_t loc_buf,
              uint32_t lkey, bool flg){
  struct ibv_sge li = {loc_buf, size, lkey};
  struct ibv_send_wr *bad_wr, send_wr;
  memset(&send_wr, 0, sizeof(send_wr));
  send_wr.wr_id = id;
  send_wr.sg_list = &li;
  send_wr.num_sge = 1;
  send_wr.opcode = IBV_WR_SEND;
  if(flg)send_wr.send_flags = IBV_SEND_SIGNALED;
  return ibv_post_send(qp, &send_wr, &bad_wr);
}
void *work_routine(void *parm){
  int ret;
  cpu_set_t cpu_mask;
  CPU_ZERO(&cpu_mask);
  CPU_SET(10, &cpu_mask);
  ret = pthread_setaffinity_np(pthread_self() ,sizeof(cpu_set_t), &cpu_mask);
  assert(ret == 0);

  ib_inf *inf = (ib_inf*)parm;
  uint64_t blk_size = inf->inp->blk_size;
  struct ibv_qp **qps = inf->local_qp;
  int num_qps = inf->inp->num_qps;
  uint64_t real_blk_size = blk_size * BUF_FACT;

  ib_mr_attr **mr_attr = (ib_mr_attr **)malloc(num_qps * sizeof(ib_mr_attr *));
  char **loc_buf = (char **)malloc(num_qps * sizeof(char *));
  ibv_mr **loc_mr = (ibv_mr **)malloc(num_qps * sizeof(ibv_mr *));
  uint64_t *laddr = (uint64_t*)malloc(num_qps * sizeof(uint64_t));
  uint64_t *raddr = (uint64_t*)malloc(num_qps * sizeof(uint64_t));
  uint32_t *lkey = (uint32_t*)malloc(num_qps * sizeof(uint32_t));
  uint32_t *rkey = (uint32_t*)malloc(num_qps * sizeof(uint32_t));

  assert(mr_attr && loc_buf && loc_mr && lkey && rkey && raddr && laddr);
  for(int i = 0; i < num_qps; i++){
    char name[STR_LEN];
    sprintf(name, "mem-%d", i);
    while(mt_get(inf->mi, name, (char **)(&mr_attr[i]), NULL))usleep(1000);

    loc_buf[i] = (char *)memalign(4*KB, real_blk_size);
    assert(loc_buf[i]);
    loc_mr[i] = ibv_reg_mr(inf->pd, loc_buf[i], real_blk_size, MR_PERMISSION);
    assert(loc_mr[i]);
    memset(loc_buf[i], '1', real_blk_size);
    laddr[i] = (uint64_t)loc_buf[i];
    raddr[i] = (uint64_t)mr_attr[i]->addr;
    lkey[i] = loc_mr[i]->lkey;
    rkey[i] = mr_attr[i]->rkey;
  }

  struct timeval t1,t2;
  ll tot_rep = inf->inp->iter_num;
  uint64_t offset;

  if(inf->test_mode == TEST_WRITE){

    gettimeofday(&t1, NULL);
    for(ll i = 1; i <= tot_rep; i++){
      offset = blk_size * ((BUF_FACT-1) & i);
      for(int j = 0; j < num_qps; j++){
        //((char *)laddr[j])[(i*(j+19260817))%blk_size] = '$'; //make a litte change?
        while(post_write(blk_size, i, qps[j], laddr[j] + offset, lkey[j],
          raddr[j] + offset, rkey[j], i % SEND_SIG_INTERVAL == 0));
      }
    }
    while(poll_flag == false);
    gettimeofday(&t2, NULL);

  }else if(inf->test_mode == TEST_READ){

    gettimeofday(&t1, NULL);
    for(ll i = 1; i <= tot_rep; i++){
      offset = blk_size * ((BUF_FACT - 1) & i);
      for(int j = 0; j < num_qps; j++){
        while(post_read(blk_size, i, qps[j], laddr[j] + offset, lkey[j],
          raddr[j] + offset, rkey[j], i % SEND_SIG_INTERVAL == 0));
      }
    }
    while(poll_flag == false);
    gettimeofday(&t2, NULL);

  }else if(inf->test_mode == TEST_SEND){

    gettimeofday(&t1, NULL);
    for(ll i = 1; i <= tot_rep; i++){
      offset = blk_size * ((BUF_FACT - 1) & i);
      for(int j = 0; j < num_qps; j++){
        while(post_send(blk_size, i, qps[j], laddr[j] + offset, 
                        lkey[j], i % SEND_SIG_INTERVAL == 0));
      }
    }
    while(poll_flag == false);
    gettimeofday(&t2, NULL);

  } else assert(0);

  ull diff_us = (ull)(t2.tv_sec - t1.tv_sec) * 1000000LL + t2.tv_usec - t1.tv_usec;
  ull tot = tot_rep * blk_size * num_qps;
  ull thgt_mb_per_s = tot/diff_us;
  double mops = (double)(tot_rep * num_qps)/diff_us;
  cout<<"work thread finish"<<endl;
  usleep(1000);

  cout<<"total us is "<<diff_us<<endl;
  cout<<"throughput is "<<thgt_mb_per_s<<"MB/s"<<endl;
  cout<<"Mops is "<<mops<<endl;
  cerr<<blk_size<<' '<<thgt_mb_per_s<<' '<<mops<<endl;
  return NULL;
}

void run_client(struct input_params *conf){
  int res;
  conf->role_int = ROLE_CLIENT;
  ib_inf *node_inf = ib_complete_setup(conf);
  
  if(strlen(conf->test_mode) == 0){
    cout<<"input the test mode :";
    res = scanf("%s",conf->test_mode);
    assert(res == 1);
  }
  if(strcmp(conf->test_mode, "write") == 0)node_inf->test_mode = TEST_WRITE;
  else if(strcmp(conf->test_mode, "read") == 0)node_inf->test_mode = TEST_READ;
  else if(strcmp(conf->test_mode, "send") == 0)node_inf->test_mode = TEST_SEND;
  else {cout<<"unkown test mode \n";return;};
  cout<<"test mode is "<<conf->test_mode<<endl;
  
  //send the mode message
  mt_set(node_inf->mi, "test_mode", (char *)&node_inf->test_mode, sizeof(node_inf->test_mode));
  mt_set(node_inf->mi, "blk_size", (char *)&conf->blk_size, sizeof(conf->blk_size));

  while(mt_get(node_inf->mi, "ready", NULL, NULL))usleep(1000);

  assert(conf->iter_num % SEND_SIG_INTERVAL == 0);

  pthread_t poll_thread, work_thread;

  poll_flag = false; //set flag
  //start the poll thread
  res = pthread_create(&poll_thread, NULL, poll_routine, node_inf);
  assert(res == 0);
  res = pthread_create(&work_thread, NULL, work_routine, node_inf);
  assert(res == 0);

  pthread_join(work_thread, NULL);
  pthread_join(poll_thread, NULL);
  mt_set(node_inf->mi, "done", "done", strlen("done"));
}

#endif